from .cann_data import *
from abc import ABC
from functools import reduce
import collections
import os
import re
class TreeTraverse:
    """
    遍历整棵树，针对所有章节和子章节进行提问题。
    """
    def __init__(self,function_handle,output_dir="./",question_depth=-1,max_conv_round = 20,fork_num = 1) -> None:
        super().__init__()
        self.function_handle = function_handle
        self.max_conv_round = max_conv_round
        self.fork_num = fork_num
        self.question_depth = question_depth
        self.output_dir = output_dir
        os.makedirs(output_dir,exist_ok=True)

    def extract_chapter_ids(self,folder_path):
        chapter_ids = []
        pattern = re.compile(r"questions_for_chapter_(\d+(?:\.\d+)*)\.json")
        
        for file_name in os.listdir(folder_path):
            match = pattern.match(file_name)
            if match:
                chapter_ids.append(match.group(1))
        
        return chapter_ids

    def visit(self,node:BaseNode):
        if isinstance(node, Section):
            self._visit_section(node)
        elif isinstance(node, Book):
            self._visit_book(node)
        else:
            raise NotImplementedError(f"No visit method defined for {node.__class__.__name__}")

    def _visit_book(self,node:Book):
        for key,child in node.children.items():
            self.visit(child)
            
    def _visit_section(self,node:Section):
        for key,child in node.children.items():
            self.visit(child)
        #最少进行6轮提问，层级越深，问问题轮数越少
        current_conv_round = max(self.max_conv_round-2*node.depth,6)
        if node.id.find("16.") is None:
            self.function_handle(root_section=node ,max_round=current_conv_round,fork_num = self.fork_num, file_name = os.path.join(self.output_dir,f"questions_for_chapter_{node.id}.json"))
        elif (self.question_depth == -1 or node.depth >= self.question_depth):
            self.function_handle(root_section=node ,max_round=current_conv_round,fork_num = self.fork_num, file_name = os.path.join(self.output_dir,f"questions_for_chapter_{node.id}.json"))

class UpdateBook:
    """
    将node的summary替换为dist的summary，两棵树必须完全一致
    """
    def __init__(self) -> None:
        super().__init__()
    def visit(self,node:BaseNode,dist:BaseNode):
        if isinstance(node, Section):
            self._visit_section(node,dist)
        elif isinstance(node, ImagePaths):
            self._visit_image_path(node,dist)
        elif isinstance(node, CodeSnippet):
            self._visit_code_snippet(node,dist)
        elif isinstance(node, Book):
            self._visit_book(node,dist)
        else:
            raise NotImplementedError(f"No visit method defined for {node.__class__.__name__}")

    def _visit_book(self,node:Book,dist:BaseNode):
        for key,child in node.children.items():
            self.visit(child,dist.children[key])
    def _visit_section(self,node:Section,dist:BaseNode):
        node.summary = dist.summary
        for key,child in node.children.items():
            self.visit(child,dist.children[key])
    def _visit_image_path(self,node:ImagePaths,dist:BaseNode):
        pass
    def _visit_code_snippet(self,node:CodeSnippet,dist:BaseNode):
        pass

class CodeGenJson():
    """
    通过树结构生成json
    """
    def __init__(self,fully_export = False) -> None:
        self.fully_export = fully_export
        super().__init__()

    def visit(self,node:BaseNode):
        if isinstance(node, Section):
            return  self._visit_section(node)
        elif isinstance(node, ImagePaths):
            return  self._visit_image_path(node)
        elif isinstance(node, CodeSnippet):
            return  self._visit_code_snippet(node)
        elif isinstance(node, Book):
            return  self._visit_book(node)
        else:
            raise NotImplementedError(f"No visit method defined for {node.__class__.__name__}")

    def _visit_book(self,node:Book):
        children = {}
        for key,child in node.children.items():
            children[key] = self.visit(child)
        return dict(sorted(children.items()))

    def _visit_section(self,node:Section):
        children = {}
        for key,child in node.children.items():
            children[key] = self.visit(child)
        code_snippet = {}
        for key,child in node.code_snippets.items():
            code_snippet[key] = self.visit(child)
        
        if self.fully_export:
            ret =  { 
                "title": node.title,
                "summary": node.summary,
                "原文内容": node.original_content,
                "图片内容": self._visit_image_path(node.image_paths),
                "代码段": dict(sorted(code_snippet.items())),
                "children": dict(sorted(children.items()))
            }
        else:
            ret =  { 
                "title": node.title,
                "summary": node.summary,
                "代码段": dict(sorted(code_snippet.items())),
                "children": dict(sorted(children.items()))
            }
        return ret

    def _visit_image_path(self,node:ImagePaths):
            return node.paths

    def _visit_code_snippet(self,node:CodeSnippet):
        return {"description":node.description,
                "codes":node.codes
        }
        
def extract_substrings(s):
    parts = s.split('.')
    results = []
    for i in range(len(parts)-1, 0, -1):
        results.append('.'.join(parts[:i]))
    return results

class MergeTreeStructure():
    """
    更新树结构，将子树merge进入主树中
    """
    def __init__(self):
        pass

    def mutate(self,origin:BaseNode,new_tree:BaseNode):
        if isinstance(origin, Section):
           return  self._mutate_section(origin,new_tree)
        elif isinstance(origin, ImagePaths):
           return     self._mutate_image_paths(origin,new_tree)
        elif isinstance(origin, CodeSnippet):
           return     self._mutate_code_snippet(origin,new_tree)
        elif isinstance(origin, Book):
           return     self._mutate_Book(origin,new_tree)
        else:
            raise NotImplementedError(f"No visit method defined for {BaseNode.__class__.__name__}")

    def mutate_dict_like_object(self,dist_like_obj_left:dict, dist_like_obj_right:dict,is_leaf=False):
        new_dict = dict(dist_like_obj_left)
        for key,value in dist_like_obj_right.items():
            if(dist_like_obj_left.get(key) is None):
                new_dict[key] = value
            else:
                new_dict[key] = self.mutate(dist_like_obj_left[key], value) if not is_leaf else value

        return new_dict  
    
    def combine_two_str(self,left:str,right:str):
        if len(left) == 0:
            return right
        if len(right) == 0:
            return left
        return left + "\n" + right

    def _mutate_section(self,origin:Section,new_tree:Section):
        new_summary = self.combine_two_str(origin.summary, new_tree.summary)
        new_original_content = self.combine_two_str(origin.original_content , new_tree.original_content)

        new_code_snippet = self.mutate_dict_like_object(origin.code_snippets, new_tree.code_snippets)
        image_paths = self.mutate(origin.image_paths,new_tree.image_paths)
        new_children = self.mutate_dict_like_object(origin.children,new_tree.children)
        return Section(
            id = origin.id,
            title = self.combine_two_str(origin.title,new_tree.title),
            summary = new_summary,
            original_content=new_original_content,
            code_snippets=new_code_snippet,
            image_paths=image_paths,
            depth = origin.depth,
            children=new_children
        )
    def _mutate_image_paths(self,origin:ImagePaths,new_tree:ImagePaths):
        new_image_path = self.mutate_dict_like_object(origin.paths,new_tree.paths,is_leaf=True)
        return ImagePaths(new_image_path)
    
    def _mutate_code_snippet(self,origin:CodeSnippet,new_tree:ImagePaths):
        return origin

    def _mutate_Book(self,origin:Book, new_tree:BaseNode):
        new_children = {}
        new_children.update(origin.children)
        for key,value in new_tree.children.items():
            if(origin.children.get(key) is None):
                new_children[key] = value
            else:
               new_children[key] = self.mutate(origin.children[key],value)
        return Book(new_children)

class CompleteTree():
    """
    生长出子树的章节root节点，之后多颗子树可以合并
    """
    def __init__(self) -> None:
        self.merger = MergeTreeStructure()
        self.outest_children = {}

    def grow_parents(self,child, parent_key, depth):
        
        return Section(
            id = parent_key,
            title="",
            summary="",
            original_content="",
            code_snippets={},
            image_paths=ImagePaths({}),
            children = {child.id:child},
            depth=depth,
        )
    def validate(self,node:BaseNode):
        if isinstance(node, Section):
          return self._validate_section(node)
        elif isinstance(node, ImagePaths):
          return  self._visit_image_path(node)
        elif isinstance(node, CodeSnippet):
           return self._validate_code_snippet(node)
        elif isinstance(node, Book):
           return self._validate_book(node)
        else:
            raise NotImplementedError(f"No visit method defined for {node.__class__.__name__}")

    def _validate_book(self,node:Book):

        for key,child in node.children.items():
            self.validate(child)

    def _validate_section(self,node:Section):
        parent_nodes = extract_substrings(node.id)
        current_depth = len(parent_nodes)
        new_children = {}
        for key, item in node.children.items():
            current_node_parents = extract_substrings(key)
            child_depth = len(parent_nodes)
            if child_depth > current_depth+1:
                # 孩子节点过深，需要外提
                self.outest_children[key] = node.children.pop(key)
                continue
            new_children[key] = self.validate(item)
            
    def _visit_image_path(self,node:ImagePaths):
        pass
    def _validate_code_snippet(self,node:CodeSnippet):
        pass
    
        
    def grow(self,origin: Book):
        self.outest_children = dict(origin.children)
        self.validate(origin)
        all_tree_roots = []
        for key,item in self.outest_children.items():
            children = {}
            parent_nodes = extract_substrings(key)
            item_depth = len(parent_nodes)
            item.depth = item_depth
            child_node = item
            parent_node = item
            for current_depth in range(0,item_depth):
                parent_node = self.grow_parents(child_node,parent_nodes[current_depth],current_depth)
                child_node = parent_node
            children[parent_node.id] = parent_node
            all_tree_roots.append(Book(children=children))
        out = reduce(self.merger.mutate,all_tree_roots, Book())
        return out

def merge_json_from_dir(path,output_path="merged_json.json"):
    """
    将多个json合并成一个大的json
    """
    chapters = read_json_files(path)
    complete = CompleteTree()
    new_chapters = []
    for current_chapter in chapters:
        chapter = complete.grow(current_chapter)
        new_chapters.append(chapter)
    merger = MergeTreeStructure()
    out = reduce(merger.mutate,new_chapters, Book())
    codegen = CodeGenJson(True)
    out_json = codegen.visit(out)
    out_json_str = json.dumps(out_json,ensure_ascii=False,indent = 4)
    # print(out_json_str)
    with open(output_path,"w") as f:
        json.dump(out_json,f,ensure_ascii=False,indent=4)
    return out_json_str
